Crank-Nicolson demo

using LinearAlgebra, Plots, Printf

This implementation is again in Julia. Since symmetric tridiagonal matrices are so prevalent and important, Julia has a special data type for them. This makes implementing Crank--Nicolson simple. If you are using python or Matlab you want to use the spdiags command. For example, in python:

import numpy as np
from scipy.sparse import spdiags

m = 4
data = np.array([np.ones(m), -2.0*np.ones(m) , np.ones(m)]);
diags = np.array([-1, 0, 1])
A = spdiags(data, diags, m, m)

A.toarray()  # just to see what it looks like

h = 0.01;
m = convert(Int64,1/h)-1;
k = 0.01;
T = 10;
A = SymTridiagonal(fill(2.0,m),fill(-1.0,m-1))
r = k/(2*h^2)
Al = I + r*A
Ar = I - r*A;

Initial and boundary data

$$ \begin{cases} u_t = u_{xx},\\ u(x,0) = \eta(x),\\ u(0,t) = g_0(t),\\ u(1,t) = g_1(t). \end{cases} $$

g0 = t -> t.^2/(1 .+ t.^2)*sin.(4*t)
g1 = t -> 0.
η = x -> exp.(-20*(x .-1/2).^2)
println(g0(0.))
println(η(0.))

0.0
0.006737946999085467

Create a gif

plot()
anim = Animation()
n = convert(Int64,ceil(T/k))
x = h:h:1-h
U = η(x)
t = 0.0
plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter)
frame(anim)

fr = 100 #frames/unit time
tb = convert(Int64,ceil(n/(fr*T)))
for i = 2:n+1
    t += k
    U = Ar*U
    U[1] += r*(g0(t)+g0(t-k))
    U[end] += r*(g1(t)+g1(t-k))
    U = Al\U
    if mod(i-1,tb) ≈ 0.0
        plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
        plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter)
        frame(anim)
    end
end
gif(anim,"heat_CN.gif")

┌ Info: Saved animation to 
│   fn = /Users/thomastrogdon/Dropbox (uwamath)/Teaching/586/amath-586-2022/notebooks/heat_CN.gif
└ @ Plots /Users/thomastrogdon/.julia/packages/Plots/PGQ1Y/src/animation.jl:114

Real-time plotting

h = 0.01;
m = convert(Int64,1/h)-1;
k = 0.01;
T = 10;
A = SymTridiagonal(fill(2.0,m),fill(-1.0,m-1))
r = k/(2*h^2)
Al = I + r*A
Ar = I - r*A;

n = convert(Int64,ceil(T/k))
x = h:h:1-h
U = η(x)
t = 0.0
plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display

fr = 100 #frames/unit time
tb = convert(Int64,ceil(n/(fr*T)))
for i = 2:n+1
    t += k
    U = Ar*U
    U[1] += r*(g0(t)+g0(t-k))
    U[2] += r*(g1(t)+g1(t-k))
    U = Al\U
    if mod(i-1,tb) ≈ 0.0
        IJulia.clear_output(true)
        plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
        plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display
    end
end

Timing without plotting

function CN_heat(η,g0,g1,T,k,h)
    m = convert(Int64,1/h)-1;
    A = SymTridiagonal(fill(2.0,m),fill(-1.0,m-1))
    r = k/(2*h^2)
    Al = I + r*A
    Ar = I - r*A;
    n = convert(Int64,ceil(T/k))
    x = h:h:1-h
    U = η(x)
    t = 0.0
    for i = 2:n+1
        t += k
        U = Ar*U
        U[1] += r*(g0(t)+g0(t-k))
        U[2] += r*(g1(t)+g1(t-k))
        U = Al\U
    end
    U
end

CN_heat (generic function with 1 method)

@time CN_heat(η,g0,g1,10.,0.0001,0.0001);

152.674556 seconds (858.32 k allocations: 29.836 GiB, 4.60% gc time, 0.03% compilation time)

Forward Euler-like discretization (what not to do...)

h = 0.01;
m = convert(Int64,1/h)-1;
k = h^2/2;  # boundary of stability region
T = 10;
A = SymTridiagonal(fill(2.0,m),fill(-1.0,m-1))
r = k/(h^2) # note that this is doubled
# Al = I + r*A # no need for this
Ar = I - r*A;

n = convert(Int64,ceil(T/k))
x = h:h:1-h
U = η(x)
t = 0.0
plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display

fr = 100 #frames/unit time
tb = convert(Int64,ceil(n/(fr*T)))
for i = 2:n+1
    t += k
    U = Ar*U
    U[1] += r*(g0(t-k))
    U[2] += r*(g1(t-k))
    # U = Al\U remove this step
    if mod(i-1,tb) ≈ 0.0
        IJulia.clear_output(true)
        plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
        plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display
    end
end

eigvals(Ar)

99-element Vector{Float64}:
 -0.9995065603657316
 -0.9980267284282716
 -0.9955619646030799
 -0.9921147013144777
 -0.9876883405951379
 -0.9822872507286885
 -0.9759167619387477
 -0.9685831611286311
 -0.9602936856769435
 -0.9510565162951536
 -0.9408807689542255
 -0.9297764858882513
 -0.9177546256839811
  ⋮
  0.9297764858882518
  0.9408807689542258
  0.9510565162951539
  0.9602936856769424
  0.9685831611286311
  0.9759167619387474
  0.9822872507286892
  0.9876883405951381
  0.9921147013144783
  0.9955619646030803
  0.9980267284282716
  0.9995065603657315

h = 0.01;
m = convert(Int64,1/h)-1;
k = h^2*.5002;  # just outside of stability region
T = 2.5;
A = SymTridiagonal(fill(2.0,m),fill(-1.0,m-1))
r = k/(h^2) # note that this is doubled
# Al = I + r*A # no need for this
Ar = I - r*A;

n = convert(Int64,ceil(T/k))
x = h:h:1-h
U = η(x)
t = 0.0
plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display

fr = 1000 #frames/unit time
tb = convert(Int64,ceil(n/(fr*T)))
for i = 2:n+1
    t += k
    U = Ar*U
    U[1] += r*(g0(t-k))
    U[2] += r*(g1(t-k))
    # U = Al\U remove this step
    if mod(i-1,tb) ≈ 0.0
        IJulia.clear_output(true)
        plot(x, U, xaxis = [0,1], yaxis = [-1,2],lw=3,label = @sprintf("u(x,t), t = %1.2f",t))
        plot!([0,1],[g0(t),g1(t)], label = "BCs", seriestype = :scatter) |> IJulia.display
    end
end